#include "types.h"
#include "params.h"
#include "memory.h"
#include "kfunc.h"
#include "lock.h"
#include "util.h"
#include "kalloc.h"
#include "memory.h"
#include "proc.h"
#include "biscuit.h"
#include "kvm.h"

extern char KERNBG[]; // kernel begin
extern char KERNED[];

char * kpgtb;
extern char trampoline[];
char * kstack[NR_PROC];
char * ktrapframe[NR_PROC];
// char * kstack;
// char * kguard;

struct {
    int valid;
    char * pgtb;
    char * va;
    char * pa;
    struct sleeplk lk;
} cache;

void
flush_cache () {
    cache.valid = 0;
    cache.pgtb = 0;
    cache.pa = 0;
    cache.va = 0;
}

void
kvminit () {
    char * adptr;
    kpgtb = kalloc ();
    char * MEMED = (char*) MAX_MEMSZ; 
    //printk ("kpgtb %x\n", (int) kpgtb);
    // ll_testnum ((int)kpgtb);
    //printk ("TRAMPOLINE :%p trampoline:%p\n", TRAMPOLINE, trampoline);
    // ll_testnum ((int)trampoline);
    for (adptr = KERNBG; adptr < MEMED; adptr += PGSZ) {
        kvmmap (adptr, adptr, kpgtb, PTE_R | PTE_W | PTE_X,0); // 将内核映射到原位置
    }
    //ll_testnum ((int)adptr);
    printk ("kernel memory mapping finished.\n");
    int i;
    for (i = 0; i < NR_PROC; i++) {
        // printk ("allocating %d\n", i);
        kstack[i] = kalloc ();
        ktrapframe[i] = kalloc ();
    }
    printk ("kstacks & trapframe memory allocation finished.\n");
    // kvmmap (KSTACK0, kstack, kpgtb, PTE_R | PTE_W);
    // kvmmap (KGUARD0, kguard, kpgtb, 0);
    // ll_testnum (TRAMPOLINE);
    //printk ("TRAMPOLINE :%p trampoline:%p\n", TRAMPOLINE, trampoline);
    
    kvmmap ((char*)TRAMPOLINE, trampoline, kpgtb, PTE_R | PTE_W | PTE_X,0); // 将trampoline映射到TRAMPOLINE
    // ((struct trapframe*) trapframe)->ksp = KSTACK0 + PGSZ;
    // ((struct trapframe*) trapframe)->kpgtb = (int) kpgtb;
    //ll_halt ();
    //panic ("before paging.");
    
    ll_setmatp (kpgtb); // 开启分页
    printk ("Virtual memory : Paging on\n");
}

/* !! 映射无需加PTE_V标志,PTE_V已在该函数中自动补齐, p == 0:内核 */
int
kvmmap (char * mapaddr, char * kpgaddr, char * pgtb,int flg, struct proc *p) {
    struct linkpgs* newpg;
    if (!PGALIGNED (mapaddr) || !PGALIGNED (kpgaddr))
        goto bad;
    pte_t * pte;
    pte = (pte_t *)pgtb; // pte = L1 pgtb
    pte += L1_PTE_INDEX (mapaddr);
    // 一级页表项地址
    if (*pte & PTE_V) { // L1 PTE有效
        pte = (pte_t *)((*pte) & ~PG4K_MASK); // pte = L2 pgtb
    } else {
        newpg = kalloc ();
        if (p)
            pushlinkpg (p->allocpgs, newpg);
        *pte = (pte_t)newpg | PTE_V;
        pte = (pte_t *)((*pte) & ~PG4K_MASK); // pte = L2 pgtb
    }
    pte += L2_PTE_INDEX (mapaddr);
    if (*pte & PTE_V) goto bad; // this address has been mapped
    *pte = (pte_t)kpgaddr | (flg & FLAG_MASK) | PTE_V; // map page
    // if ((unsigned int)mapaddr == 0xFFF7F000) {
    //     printk ("0xFFF7F000 has been mapped.\n");
    // }

    return 0;
    bad:
    return -1;
}

/* va 转 pa (4K aligned), va 4K aligned is faster */
int
kvmv2p4k (char * pgtb, char * vaddr, char **paddr) {
    sleep_lock (&cache.lk);
    if (cache.valid)
    if (cache.pgtb == pgtb)
    if (cache.va == vaddr) {
        *paddr = cache.pa;
        // printk ("pa:%p\n", cache.pa);
        // panic ("kvmv2p4k");
        sleep_unlock (&cache.lk);
        goto paloaded;
    }
    pte_t * pte;
    pte = (pte_t *) pgtb; // pte = L1 pgtb
    // goto L1 pte
    pte += L1_PTE_INDEX (vaddr);
    if (*pte & PTE_V) { // L1 PTE有效
        pte = (pte_t *)((*pte) & ~PG4K_MASK); // pte = L2 pgtb
    } else {
        goto bad;
    }
    // now pte is pointing to L2 pgtb
    pte += L2_PTE_INDEX (vaddr);
    if (*pte & PTE_V) {
       pte = (pte_t *)((*pte) & ~PGMASK); 
    } else {
        goto bad;
    }
    // now pte is pointing to aligned phyaddr
    //sleep_lock (&cache.lk);
    cache.va = vaddr;
    cache.pa = (char *)pte;
    cache.pgtb = pgtb;
    cache.valid = 1;
    *paddr = (char *)pte;
    sleep_unlock (&cache.lk);
    paloaded:
    return 0;
    bad:
    sleep_unlock (&cache.lk);
    return -1;
}

int
vmread (char * pgtb, unsigned int from, char * to, unsigned int sz) {
    // unsigned int va;
    char * pa, * va;
    unsigned int foff;
    va = (char *) from;
    while (sz) {
        foff = (unsigned int)va & PG4K_MASK; //? va & PG4K_MASK : 0x1000;
        if (kvmv2p4k (pgtb, va, &pa))
            goto error;
        if (va == 0)
            printk ("va:%p, pa:%p\n",va,pa);
        for (; foff < 0x1000; foff++) {
            *(char*)to = *(char*)(pa + foff);
            to++;
            sz--;
            if (sz == 0)
                break;
        }
        va = (char*)((unsigned int)va & ~PG4K_MASK) + PGSZ;
    }
    return 0;
    error:
    return -1;
}

int
vmwrite (char * pgtb, char * from , unsigned int to, unsigned int sz) {
    // unsigned int va;
    char * pa, * va;
    unsigned int toff;
    va = (char*) to;
    while (sz) {
        toff = (unsigned int)va & PG4K_MASK; //? va & PG4K_MASK : 0x1000;
        if (kvmv2p4k (pgtb, va, &pa))
            goto error;
        for (; toff < 0x1000; toff++) {
            *(char*)(pa + toff) = *(char*)from;
            from++;
            sz--;
            if (sz == 0)
                break;
        }
        va = (char*)((unsigned int)va & ~PG4K_MASK) + PGSZ;
    }
    return 0;
    error:
    return -1;
}

